#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.
import os
from pathlib import Path
import json
import numpy as np
import tqdm
from PIL import Image
import cv2
from panopticapi.utils import IdGenerator, save_json
import pycocotools.mask as mask_util
import glob

data_dir = Path('/root/autodl-tmp/data/zzdx_202407/')

def handle_train_files():
    src_dir = data_dir / "raw_train"
    src_imgs_dir = src_dir / "images"
    src_masks_dir = src_dir / "masks"
    
    dest_dir = data_dir / "train"
    
    index = 0
    for file in tqdm.tqdm(list(src_masks_dir.iterdir())):
        farr = file.name.split("@")
        # if farr[-5] != "-60":
        #     continue
        
        out_img_dir = dest_dir / "images" 
        out_mask_dir = dest_dir / "masks" 
        out_img_dir.mkdir(parents=True, exist_ok=True)
        out_mask_dir.mkdir(parents=True, exist_ok=True)
        
        in_img_file = src_imgs_dir / file.name
        in_mask_file = src_masks_dir / file.name
        out_img_file = out_img_dir / file.name
        out_mask_file = out_mask_dir / file.name
        
        convert(in_img_file,in_mask_file,out_img_file,out_mask_file)
        
def handle_valsim_files():
    src_dir = data_dir / "raw_val_sim"
    src_imgs_dir = src_dir / "images"
    src_masks_dir = src_dir / "masks"
    
    dest_dir = data_dir /"processed_dir"/ "val_sim"
    
    index = 0
    for file in tqdm.tqdm(list(src_masks_dir.iterdir())):
        farr = file.name.split("@")
        index += 1
        if index % 3 != 0:
            continue
        
        out_img_dir = dest_dir / "images" 
        out_mask_dir = dest_dir / "masks" 
        out_img_dir.mkdir(parents=True, exist_ok=True)
        out_mask_dir.mkdir(parents=True, exist_ok=True)
        
        in_img_file = src_imgs_dir / file.name
        in_mask_file = src_masks_dir / file.name
        out_img_file = out_img_dir / file.name
        out_mask_file = out_mask_dir / file.name
        
        convert_2(in_img_file,in_mask_file,out_img_file,out_mask_file)
        
def convert_2(in_img_file,in_mask_file, out_img_file,out_mask_file):
    in_img = Image.open(in_img_file)
    in_mask = Image.open(in_mask_file)
    w2 = int(in_img.size[0]/2)
    h2 = int(in_img.size[1]/2)
    in_img = in_img.resize((w2, h2), resample=Image.BILINEAR)
    in_mask = in_mask.resize((w2, h2), resample=Image.NEAREST)
    in_img = np.array(in_img)
    in_mask = np.array(in_mask)
    assert in_img.dtype == np.uint8
    assert in_mask.dtype == np.uint8
    in_mask.flags.writeable = True
    in_mask[:,:,1] = in_mask[:,:,0] > 0
    in_mask[:,:,[0,1]] = in_mask[:,:,[1,0]]
    in_mask = np.delete(in_mask,3,axis=2)
    in_img = np.delete(in_img,3,axis=2)
    
    ids = np.unique(in_mask[:,:,1])
    if len(ids) < 4:
        return
    
    Image.fromarray(in_img).save(out_img_file)
    Image.fromarray(in_mask).save(out_mask_file)
   
def convert(in_img_file,in_mask_file, out_img_file,out_mask_file):
    in_img = Image.open(in_img_file)
    in_mask = Image.open(in_mask_file)
    in_img = np.array(in_img)
    in_mask = np.array(in_mask)
    assert in_img.dtype == np.uint8
    assert in_mask.dtype == np.uint8
    in_mask.flags.writeable = True
    in_mask[:,:,1] = in_mask[:,:,0] > 0
    in_mask[:,:,[0,1]] = in_mask[:,:,[1,0]]
    in_mask = np.delete(in_mask,3,axis=2)
    in_img = np.delete(in_img,3,axis=2)
    
    ids = np.unique(in_mask[:,:,1])
    if len(ids) < 2:
        return
    
    Image.fromarray(in_img).save(out_img_file)
    Image.fromarray(in_mask).save(out_mask_file)
    

def convert_val(in_img_file,in_mask_file, out_img_file,out_mask_file):
    if not in_mask_file.exists():
        return
    in_img = Image.open(in_img_file)
    in_mask = Image.open(in_mask_file)
    w2 = int(in_img.size[0]/2)
    h2 = int(in_img.size[1]/2)
    in_img = in_img.resize((w2, h2), resample=Image.BILINEAR)
    in_mask = in_mask.resize((w2, h2), resample=Image.NEAREST)
    in_img = np.array(in_img)
    in_mask = np.array(in_mask)
    assert in_img.dtype == np.uint8
    assert in_mask.dtype == np.uint8
    w4 = w2 //2 
    in_img = in_img[:,w4-256: w4+256,:]
    in_mask = in_mask[:,w4-256: w4+256,:]
    in_mask.flags.writeable = True
    in_mask[:,:,1] = in_mask[:,:,0] > 0
    in_mask[:,:,[0,1]] = in_mask[:,:,[1,0]]
    # in_mask = np.delete(in_mask,3,axis=2)
    # in_img = np.delete(in_img,3,axis=2)

    ids = np.unique(in_mask[:,:,1])
    if len(ids) < 4:
        return
    
    Image.fromarray(in_img).save(out_img_file)
    Image.fromarray(in_mask).save(out_mask_file)

def handle_raw_val_files():
    
    see_files = Path("/root/autodl-tmp/data/zzdx_202407/raw_val_real/masks")
    
    src_dir = data_dir / "raw_val_real"
    src_imgs_dir = src_dir / "images"
    src_masks_dir = src_dir / "masks"
    
    dest_dir = data_dir / "processed_dir" / "val_real"
    
    index = 0
    for file in tqdm.tqdm(list(see_files.iterdir())):
        fname = file.name
        index += 1
        if index % 3 != 0:
            continue
        out_img_dir = dest_dir / "images" 
        out_mask_dir = dest_dir / "masks" 
        out_img_dir.mkdir(parents=True, exist_ok=True)
        out_mask_dir.mkdir(parents=True, exist_ok=True)
        
        in_img_file = src_imgs_dir / file.name.replace(".png",".jpg")
        in_mask_file = src_masks_dir / file.name
        out_img_file = out_img_dir / file.name
        out_mask_file = out_mask_dir / file.name
        
        convert_val(in_img_file,in_mask_file,out_img_file,out_mask_file)

### 第一步，生成MASK格式图像和train/val目录结构
def gen_mask():
    raw_dir = data_dir / "raw"
    processed_dir = data_dir / "processed_dir"
    train_dir = processed_dir / "train"
    val_dir = processed_dir / "val"  
    processed_dir.mkdir(parents=True, exist_ok=True)
    train_dir.mkdir(parents=True, exist_ok=True)
    val_dir.mkdir(parents=True, exist_ok=True)
    raw_img_dir = raw_dir / "SimImages"
    raw_mask_dir = raw_dir / "SimImagesMask"
    
    index = 0
    for file in tqdm.tqdm(list(raw_img_dir.iterdir())):
        output_dir = train_dir
        if index % 5 == 0:
            output_dir = val_dir  
        index += 1
        in_img_file = raw_img_dir / file.name
        in_mask_file = raw_mask_dir / file.name
        out_img_dir = output_dir / "images" 
        out_mask_dir = output_dir / "masks" 
        out_img_dir.mkdir(parents=True, exist_ok=True)
        out_mask_dir.mkdir(parents=True, exist_ok=True)
        out_img_file = out_img_dir / file.name
        out_mask_file = out_mask_dir / file.name
        convert(in_img_file,in_mask_file,out_img_file,out_mask_file)


### 第二步  生成 ade20k_instance_imgCatIds.json 和 ade20k_instance_catid_mapping.txt
def gen_ins_catids():
    
    image_id = 0
    images = []
    for name, dirname in [("train", "train"), ("val_sim", "val_sim"), ("val_real", "val_real")]:
        processed_dir = data_dir / "processed_dir"
        tv_dir = processed_dir / dirname / "images"
        for file in tqdm.tqdm(list(tv_dir.iterdir())):
            image = {}
            image["file_name"] = file.name
            image["id"] = image_id
            image_id += 1
            image["width"] = 512
            image["height"] = 512
            images.append(image)
    
    categories = []
    cat = {}
    cat["id"] = 1
    cat["name"] = "building"
    categories.append(cat)
    
    instance_imgCatIds = {}
    instance_imgCatIds["images"] = images
    instance_imgCatIds["categories"] = categories
    
    instance_imgCatIds_file = processed_dir / "instance_imgCatIds.json"
    save_json(instance_imgCatIds,instance_imgCatIds_file)
    
    with open(processed_dir / "instance_catid_mapping.txt",'w') as f2:
        f2.write("Instacne100\tSceneParse150\tFullADE20K\n")
        f2.write("1\t1\t311")
        f2.flush()
    pass

##第三步  生成COCO格式JSON
def gen_coco_json():
    processed_dir = data_dir/ "processed_dir"
    
    insid_map_cid = {}
    
    for name, dirname in [("train", "train"),  ("val_sim", "val_sim"), ("val_real", "val_real")]:
    # for name, dirname in [("val_sim", "val_sim"), ("val_real", "val_real")]:


        image_dir = processed_dir/ dirname / "images"
        instance_dir = processed_dir/ dirname / "masks"

        # img_id = 0
        ann_id = 1

        # json
        # out_file = os.path.join(str(processed_dir), f"/instance_{name}.json")
        out_file = processed_dir/ f"instance_{name}.json"

        # json config
        instance_config_file = processed_dir / "instance_imgCatIds.json"
        with open(instance_config_file) as f:
            category_dict = json.load(f)["categories"]

        # load catid mapping
        # it is important to share category id for both instance and panoptic annotations
        mapping_file = processed_dir / "instance_catid_mapping.txt"
        with open(mapping_file) as f:
            map_id = {}
            for i, line in enumerate(f.readlines()):
                if i == 0:
                    continue
                ins_id, sem_id, _ = line.strip().split()
                # shift id by 1 because we want it to start from 0!
                # ignore_label becomes 255
                map_id[int(ins_id)] = int(sem_id)

        for cat in category_dict:
            cat["id"] = map_id[cat["id"]]

        filenames = sorted(glob.glob(os.path.join(image_dir, "*.png")))

        ann_dict = {}
        images = []
        annotations = []

        for idx, filename in enumerate(tqdm.tqdm(filenames)):
            image = {}
            image_id = os.path.basename(filename)[:-4]
            if name == "train":
                scene_id = int(filename.split('@')[-3])
            else:
                scene_id = 1

            image["id"] = image_id
            image["file_name"] = os.path.basename(filename)

            original_format = np.array(Image.open(filename))
            image["width"] = original_format.shape[1]
            image["height"] = original_format.shape[0]

            images.append(image)

            filename_instance = os.path.join(str(instance_dir), image_id + ".png")
            ins_seg = np.asarray(Image.open(filename_instance))
            assert ins_seg.dtype == np.uint8

            # 标签R波段表示类别
            # GB波段表示ID
            # B*256 + G
            # because 0 is reserved as VOID label
            ##
            instance_cat_ids = ins_seg[..., 0]
            bb = np.array(ins_seg[...,2],dtype=np.int32) << 8
            gg = ins_seg[..., 1]
            instance_ins_ids = gg + bb
            #np.array(ins_seg[...,2],dtype=np.int32) << 8 
            # process things
            for thing_id in np.unique(instance_ins_ids):
                if thing_id == 0:
                    continue
                mask = instance_ins_ids == thing_id
                instance_cat_id = np.unique(instance_cat_ids[mask])
                assert len(instance_cat_id) == 1

                anno = {}
                anno['id'] = ann_id
                anno['ins_id_in_scene'] = int(thing_id)
                rsod =  scene_id * 100000 + int(thing_id)
                anno['ins_id_raw'] = rsod
                if rsod not in insid_map_cid:
                    insid_map_cid[rsod] = len(insid_map_cid.items())+1 
                anno['ins_id'] = insid_map_cid[rsod]
                
                ann_id += 1
                anno['image_id'] = image['id']
                anno["iscrowd"] = int(0)
                anno["category_id"] = int(map_id[instance_cat_id[0]])

                inds = np.nonzero(mask)
                ymin, ymax = inds[0].min(), inds[0].max()
                xmin, xmax = inds[1].min(), inds[1].max()
                anno["bbox"] = [int(xmin), int(ymin), int(xmax - xmin + 1), int(ymax - ymin + 1)]
                # if xmax <= xmin or ymax <= ymin:
                #     continue
                rle = mask_util.encode(np.array(mask[:, :, None], order="F", dtype="uint8"))[0]
                rle["counts"] = rle["counts"].decode("utf-8")
                anno["segmentation"] = rle
                anno["area"] = int(mask_util.area(rle))
                annotations.append(anno)

        # save this
        ann_dict['images'] = images
        ann_dict['categories'] = category_dict
        ann_dict['annotations'] = annotations

        save_json(ann_dict, out_file)


if __name__ == "__main__":
    # handle_train_files()
    # handle_raw_val_files()
    # handle_valsim_files()
    # gen_ins_catids()
    gen_coco_json()
